﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Runtime.InteropServices;

namespace LPLib
{
    /// <summary>
    /// Wraps the phylip DNAML routine.
    /// </summary>
    public static class DNAMLWrapper
    {
        /*
           Default Settings
             U                 Search for best tree?  Yes
             T        Transition/transversion ratio:  2.0000
             F       Use empirical base frequencies?  Yes
             C                One category of sites?  Yes
             R           Rate variation among sites?  constant rate
             W                       Sites weighted?  No
             S        Speedier but rougher analysis?  Yes
             G                Global rearrangements?  No
             J   Randomize input order of sequences?  No. Use input order
             O                        Outgroup root?  No, use as outgroup species  1
             M           Analyze multiple data sets?  No
             I          Input sequences interleaved?  Yes
           */

        /* settings (options) usage:
         * 1. only set an option if non default is needed. 
         * 1. options are separated by ;
         * 2. if an item expects true or false, then just give the option item. such as U;
         * 3. if an item expects values, then use, e.g. F=0.3 0.3 0.3 0.3;
         *      (the values are separated by space)
         * 4. for F(frequence), give 4 values, such as F=0.3 0.2 0.3 0.2 (freq a c g t)
         * 5. for C(one category persite), give two sets of values as, e.g.
         *      C=5 0.2 0.3 0.4 0.5 0.5;
         *      the first parameter gives number of categories (5)
         *      then list the rate for each category
         *      
         *      For C options, also need to set the category variable, which is a string 
         *      containing one char for each site ( 1-9 A-Z, maximum 35).
         *      
         * 6. for R (rate variation), 
         *    default   - constant rate - no inuput needed.
         *    R1=x y;    - Gamma distributed rate; 
         *              x= Coefficient of variation of substitution rate among sites (must be positive)
         *                 In gamma distribution parameters, this is 1/(square root of alpha)
         *              y= Rates in HMM: Number of categories (1-9)?
         *    R2=x y z - Gamma + Invariant Sites
         *              x same as above
         *              y same as above
         *              z Fraction of invariant sites?
         *    R3=n r1 r2...rn p1 p2...pn -user-defined HMM of rates
         *              n number of categories
         *              r1 r2...rn = rate for each category
         *              p1 p2...pn = probability for each category, adds up to 1
         *             
         *    for the above 3 options, may also optionally set A 
         *    A - default: Rates at adjacent sites correlated?  No, they are independent
         *    or 
         *    A=x; (x: Mean block length of sites having the same rate (greater than 1))
         *    
         * 7  for G, it is a switch. just use G;
         * 
         * 8. J=x y;  x: random number seed, y: number of times to jumble.
         * 9. O=x;    x: number of outgroup
         * 10. M=W/D Datasets
         */

        [DllImport("dnaml.dll", CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        private static extern int run_dnaml(string options, string weights, string categories, string input_data,
    StringBuilder output_tree, int tree_buf_len, StringBuilder outfile_content, int content_buf_len);

        //        private static extern int run_dnaml(string options, string weights, string categories, string input_data, StringBuilder output_tree, int tree_buf_len);

        ///// <summary>
        ///// Infers the maximum likelihod tree using Phylip's DNAML.
        ///// </summary>
        ///// <param name="pc">The sequence collection to be analyzed.</param>
        ///// <param name="rootID">The position of the sequence to be taken as the root.</param>
        ///// <returns>The newick string of the maximum likelihood tree.</returns>
        //public static string RunDNAML(PolymerCollection<byte> pc, string rootName)
        //{
        //    string newick = string.Empty;
        //    Dictionary<string, string> nameKey = getSafeNames(pc);
        //    string infile = WritePhylip(pc);
        //    int rootID = 0;
        //    foreach (Polymer<byte> p in pc)
        //    {
        //        rootID++;
        //        if (nameKey[p.Name] == rootName)
        //        {
        //            break;
        //        }
        //    }
        //    // TODO/DEBUG do we need the capacity here?
        //    StringBuilder outputTree = new StringBuilder(capacity: 20000);
        //    StringBuilder outfile = new StringBuilder(capacity: 60000);
        //    string options = string.Empty;
        //    if (rootID <= pc.Count)
        //        options = "O=" + rootID.ToString() + ";";
   
        //    int status = run_dnaml(options, string.Empty, string.Empty, infile, outputTree, outputTree.Capacity, outfile, outfile.Capacity);

        //    if (status == 0)
        //    {
        //        newick = CleanNewick(outputTree.ToString());
        //        Tree<int> tree = new Tree<int>(newick);
        //        RestoreNodeNames(tree, nameKey);
        //        return tree.GetDescriptor();
        //    }
        //    else
        //    {
        //        return string.Empty;
        //    }
        //}

        ///// <summary>
        ///// Infers the maximum-likelihood phylogenetic tree using DNAML.
        ///// </summary>
        ///// <param name="pc">The Polymer collection to be analyzed. Must be aligned.</param>
        ///// <param name="options">The user options and parameters for DNAML.</param>
        ///// <param name="categories">A string providing the category indentification by column in the alignment.</param>
        ///// <param name="weights"></param>
        ///// <returns>The maximum likelihood tree in Newick format.</returns>
        //public static string[] RunDNAML(PolymerCollection<byte> pc, string options, string categories = null, string weights = null)
        //{
        //    Dictionary<string, string> nameKey = getSafeNames(pc);
        //    string infile = WritePhylip(pc);

        //    StringBuilder outputTree = new StringBuilder(capacity: 20000);
        //    StringBuilder outputInformation = new StringBuilder(capacity: 60000);

                    
        //    int status = run_dnaml(options, weights, categories, infile, outputTree, outputTree.Capacity, outputInformation, outputInformation.Capacity);

        //    string newick = string.Empty;
        //    if (status == 0)
        //    {
        //        newick = CleanNewick(outputTree.ToString());
        //        Tree<int> tree = new Tree<int>(newick);
        //        RestoreNodeNames(tree, nameKey);
        //        return new string[] { tree.GetDescriptor(), outputInformation.ToString() };
        //    }
        //    else
        //    {
        //        return new string[] { string.Empty, outputInformation.ToString() };
        //    }
        //}

        /// <summary>
        /// Infers the maximum-likelihood phylogenetic tree using DNAML.
        /// </summary>
        /// <param name="pc">The Polymer collection to be analyzed. Must be aligned.</param>
        /// <param name="options">The user options and parameters for DNAML.</param>
        /// <param name="categories">A string providing the category indentification by column in the alignment.</param>
        /// <param name="weights"></param>
        /// <returns>The maximum likelihood tree in Newick format.</returns>
        public static string[] RunDNAML(PolymerCollection<byte> pc, string rootName = "", double[] categories = null, int gammaCategories = 1, double cv = 1, string categoryMap = null, string weights = null)
        {
            Dictionary<string, string> nameKey = getSafeNames(pc);
            string infile = WritePhylip(pc);

            string options = string.Empty;
            if (rootName.Length > 0)
            {
                int rootIndex = 0;
                foreach (Polymer<byte> p in pc)
                {
                    rootIndex++;
                    if (nameKey[p.Name] == rootName)
                    {
                        break;
                    }
                }
                options += "O=" + rootIndex.ToString() + ";";
            }

            if (categories != null && categories.Length > 1)
            {
                options += "C=" + categories.Length.ToString();
                foreach (double c in categories)
                {
                    options += " " + c.ToString();
                }
                options += ";";
            }

            if (gammaCategories > 1)
            {
                options += "R1=" + cv.ToString() + " " + gammaCategories.ToString() + ";";
            }

            pc.Rename(nameKey);

            StringBuilder outputTree = new StringBuilder(capacity: 20000);
            StringBuilder outputInformation = new StringBuilder(capacity: 60000);

            int status = run_dnaml(options, weights, categoryMap, infile, outputTree, outputTree.Capacity, outputInformation, outputInformation.Capacity);

            string newick = string.Empty;
            if (status == 0)
            {
                newick = CleanNewick(outputTree.ToString());
                Tree<int> tree = new Tree<int>(newick);
                RestoreNodeNames(tree, nameKey);
                return new string[] { tree.GetDescriptor(), outputInformation.ToString() };
            }
            else
            {
                return new string[] { string.Empty, outputInformation.ToString() };
            }

        }

        /// <summary>
        /// Strips out end of line characters.
        /// </summary>
        /// <param name="newick">The input string.</param>
        /// <returns>The output string = the input string with all \r and \n characters removed.</returns>
        public static string CleanNewick(string newick)
        {
            // make the necessary changes to newick
            // first, strip out the control  characters
            newick = newick.Replace("\r", string.Empty);
            newick = newick.Replace("\n", string.Empty);

            return newick;
        }

        /// <summary>
        /// Reorganizes the tree to have a root. 
        /// </summary>
        /// <param name="newick">The input tree in Newick format.</param>
        /// <returns>The rerooted tree in Newick format.</returns>
        public static string RerootNewick(string newick)
        {
            int lastCommaPosition = newick.LastIndexOf(',');

            newick = "(" + newick.Substring(0, lastCommaPosition) + ")" + newick.Substring(lastCommaPosition + 1, newick.Length - lastCommaPosition - 1);

            newick = newick.Replace("UCA", string.Empty);
            newick = newick.Substring(0, newick.Length - 1) + "UCA:0;";

            return newick;
        }

        //private static Dictionary<string, string> getSafeNames(PolymerCollection<byte> pc, ref string options)
        //{
        //    // if one of the sequences is to be used as the root, find its name
        //    string rootName = string.Empty;
        //    int iPosition = -1;
        //    string firstHalf = options;
        //    string secondHalf = string.Empty;
        //    List<string> names;
        //    bool needRoot = false;
        //    if (options.Contains("O="))
        //    {
        //        needRoot = true;
        //        names = pc.GetNames();
        //        iPosition = options.IndexOf("O=");
        //        firstHalf = options.Substring(0, iPosition);
        //        secondHalf = options.Substring(iPosition + 2);
        //        int i2 = secondHalf.IndexOf(';');
        //        int rootIndex = int.Parse(secondHalf.Substring(0, i2)) - 1;
        //        secondHalf = secondHalf.Substring(i2);
        //        rootName = names[rootIndex];
        //    }

        //    Dictionary<string, string> originalNames = new Dictionary<string, string>();
        //    Dictionary<string, string> newNames = new Dictionary<string, string>();
        //    int i = 0;
        //    foreach (Polymer<byte> p in pc)
        //    {
        //        // if the name is UCA, skip it.
        //        if (p.Name == "UCA")
        //        {
        //            originalNames.Add("UCA", "UCA");
        //            continue;
        //        }

        //        originalNames.Add(i.ToString(), p.Name);
        //        newNames.Add(p.Name, i.ToString());
        //        i++;
        //    }
        //    pc.Rename(newNames);

        //    if (needRoot)
        //    {
        //        names = pc.GetNames();
        //        for (int j = 0; j < names.Count; j++)
        //        {
        //            if (names[j] == newNames[rootName])
        //            {
        //                options = firstHalf + "O=" + (j+1).ToString() + secondHalf;
        //                break;
        //            }
        //        }
        //    }

        //    return originalNames;
        //}

        private static Dictionary<string, string> getSafeNames(PolymerCollection<byte> pc)
        {
            Dictionary<string, string> originalNames = new Dictionary<string, string>();
            Dictionary<string, string> newNames = new Dictionary<string, string>();
            int i = 0;
            foreach (Polymer<byte> p in pc)
            {
                originalNames.Add(i.ToString(), p.Name);
                newNames.Add(p.Name, i.ToString());
                i++;
            }
            pc.Rename(newNames);
            return originalNames;
        }

        private static void RestoreNodeNames(Tree<int> tree, Dictionary<string, string> idkey)
        {
            if (idkey.ContainsKey(tree.Name))
            {
                tree.Name = idkey[tree.Name];
            }

            foreach (Tree<int> child in tree.Children)
            {
                RestoreNodeNames(child, idkey);
            }
        }

        public static Tuple<List<double>, List<double>, List<int>> ParseDNAMLOutput(string output)
        {
            // read the output string and extract the discrete gamma parameters, 
            // the probabilities for each, and the integer array of indices

            List<double> rates = new List<double>();
            List<double> probabilities = new List<double>();
            List<int> indices = new List<int>();

            string[] parsed = output.Split('\n');

            // search parsed for the component that begins "State"
            // skip one component, then read while component is non-blank lines of the form
            // integer double double, representing index rate probability
            for (int i = 0; i < parsed.Length; i++)
            {
                if (parsed[i].Contains("State"))
                {
                    i += 2;
                    while (parsed[i] != "")
                    {
                        string[] components = parsed[i].Split(' ');
                        int k = 0;
                        for (int j = 0; j < components.Length; j++)
                        {
                            if (components[j] != "")
                            {
                                if (k == 1)
                                {
                                    rates.Add(double.Parse(components[j]));
                                }
                                else if (k == 2)
                                {
                                    probabilities.Add(double.Parse(components[j]));
                                }
                                k++;
                            }
                        }
                        i++;
                    }
                }
                // then search for component beginning "Combination"
                // skip one component, then read while component is non-blank
                // each line is a string of single digits, skip blanks
                else if (parsed[i].Contains("Combination"))
                {
                    i += 2;
                    while (parsed[i] != "")
                    {
                        for (int j = 0; j < parsed[i].Length; j++)
                        {
                            if (parsed[i][j] != ' ')
                            {
                                indices.Add(int.Parse(parsed[i][j].ToString()) - 1); // subtract one for zero-based indexing
                            }
                        }
                        i++;
                    }
                }
            }

            return new Tuple<List<double>, List<double>, List<int>>(rates, probabilities, indices);
        }

        /// <summary>
        /// Returns a string in Phylip format.
        /// </summary>
        /// <param name="pc">The Polymer collection to be serialized.</param>
        /// <returns>A string containing the Phylip formatted representation of the sequences in the Polymer collection.</returns>
        public static string WritePhylip(PolymerCollection<byte> pc)
        {
            StringBuilder phyString = new StringBuilder();
            int length = pc.MaxLength;
            phyString.Append(pc.Count + " " + length + "\r\n");

            foreach (Polymer<byte> p in pc)
            {
                int nameLength = p.Name.Length;

                phyString.Append(p.Name.Substring(0, Math.Min(10, nameLength)));
                for (int i = 0; i < 10 - nameLength; i++)
                    phyString.Append(" ");

                phyString.Append(new string(TypeConverter.ToChars(p.Seq)) + "\r\n");
            }
            return phyString.ToString();
        }
    }
}
